int prune(__local int *twig, __local int *base, int N, int baseSize, int twigSize)
{
	int groupCount = 0;
	int i, j;
	int idxi, idxj;

  for(i = 0; i < baseSize; i++) {
    idxi = 2 * i;
    for(j = 0; j < twigSize; j++) {
      idxj = 2 * j;
      if ((base[idxi] == twig[idxj]) || (base[idxi+1] == twig[idxj]))
        twig[idxj] = N + groupCount;
      if ((base[idxi] == twig[idxj+1]) || (base[idxi+1] == twig[idxj+1]))
        twig[idxj+1] = N + groupCount;
    }
    groupCount++;
  }
  
  for(i = 0; i < twigSize; i++) {
    idxi = 2 * i;
    if (twig[idxi] == twig[idxi + 1])
      twig[idxi] = -1; // This should happen a lot
    else {
      for(j = i + 1; j < twigSize; j++) {
        idxj = 2 * j;
        if ((twig[idxi] == twig[idxj]) || (twig[idxi+1] == twig[idxj]))
          twig[idxj] = N + groupCount;
        if ((twig[idxi] == twig[idxj+1]) || (twig[idxi+1] == twig[idxj+1]))
          twig[idxj+1] = N + groupCount;
      }
      groupCount++;
    }
  }
		
	return groupCount - baseSize;
}

__kernel void dendrogram(
	__global int *sorted,
	__global int *result,
  __global int *baseStart,
  __global int baseStartSize,
	__global int N,
	__global int branchSize,
	__global int twigSize,
  __local int *base,
  __local int *baseClean,
	__local int *branch,
	__local int *branchClean,
	__local int *count)
{
		int wg_id = get_group_id(0); 
		int th_id = get_local_id(0);
		int numThreads = get_local_size(0);
		int numGroups = get_num_groups(0);
    int th_index = th_id * twigSize * 2;
    int wg_index = wg_id * branchSize * 2;
    int result_index = wg_id * (N - baseStartSize - 1) * 2;
		int i, j;
	  
    int totalSize = 0;
    int baseSize = baseStartSize;
    int pruneSize = 0;
    int bookmark = 0;
    int bookmarkStart = 0;
    int offset = 0;
    
    int spin = 0;
    
//    for (i = 0; i < baseSize * 2; i += numThreads) //Copy in base
//      if(i + th_id < baseSize * 2)
//        base[i + th_id] = baseStart[i + th_id];

    if(th_id == 0)
      for (i = 0; i < baseSize * 2; i++) //Copy in base
        base[i] = baseStart[i];

        
//    barrier(CLK_LOCAL_MEM_FENCE);
//        
//    if(th_id == 0)
//        for(i = 0; i < 20 * 2; i++)
//          result[result_index + i] = base[i];
    
    while(branchSize - bookmark + pruneSize >= numThreads * twigSize && spin < 10 && baseSize < (N - 1)) {
      
      bookmarkStart = (bookmark - pruneSize) * 2;
      
      for (i = 0; i < numThreads * twigSize * 2; i+=numThreads) // Copy in branch
        if (i + th_id < pruneSize * 2)
          branchClean[i + th_id] = branch[i + th_id] = sorted[wg_index + i + th_id];
        else if (i + th_id < branchSize)
          branchClean[i + th_id] = branch[i + th_id] = sorted[wg_index + bookmarkStart + i + th_id];

      
      barrier(CLK_LOCAL_MEM_FENCE);

      count[th_id] = prune(&branch[th_index], base, N, baseSize, twigSize);
      
      barrier(CLK_LOCAL_MEM_FENCE);
                
      offset = 0;
      for(i = 0; i < th_id; i++)
        offset += count[i];
      offset = offset - count[0]; // Starting loop at 1 causes trouble
        
      j = 0;
      for (i = 0; i < twigSize; i++)
        if(branch[th_index + i * 2] != -1) {
          if(th_id != 0) {
            sorted[wg_index + (offset + j) * 2] = branchClean[th_index + i * 2];
            sorted[wg_index + (offset + j) * 2 + 1] = branchClean[th_index + i * 2 + 1];
          } else {
            base[(baseSize + j) * 2] = branch[i * 2]; // Copy new base (with group #s)
            base[(baseSize + j) * 2 + 1] = branch[i * 2 + 1];
            baseClean[(baseSize - baseStartSize + j) * 2] = branchClean[i * 2]; // Copy new base (without group #s)
            baseClean[(baseSize - baseStartSize + j) * 2 + 1] = branchClean[i * 2 + 1];
          }
          j++;
        }
     
      baseSize += count[0];

      //result[spin] = count[0];
              
      barrier(CLK_LOCAL_MEM_FENCE);        

      bookmark += numThreads * twigSize - pruneSize; // Move bookmark forward by the number of elements we read off
      pruneSize = 0;
      for(i = 1; i < numThreads; i++)
        pruneSize += count[i];
        
      spin++;
    
      barrier(CLK_LOCAL_MEM_FENCE);
    }
    
//      if(th_id == 0)
//        for(i = 0; i < baseSize * 2; i++)
//          result[i] = baseClean[i];

//    int totalBase = (baseSize - baseStartSize) * 2;
//    for (i = 0; i < totalBase; i+=numThreads) // Copy out base to result
//      if(i + th_id < totalBase)
//        result[result_index + i + th_id] = baseClean[i + th_id];

      if(th_id == 0)
        for(i = 0; i < (baseSize - baseStartSize) * 2; i++)
          result[result_index + i] = baseClean[i];

    if(baseSize < N - 1) { // This shouldn't happen very often if the branches are big
      for (i = 0; i < pruneSize * 2; i+=numThreads) //Copy in prune
        if(i + th_id < pruneSize * 2)
          branchClean[i + th_id] = branch[i + th_id] = sorted[wg_index + i + th_id];
    
      for(i = 0; i < (branchSize - bookmark) * 2; i+=numThreads) // Copy over remainder
        if(i + th_id < (branchSize - bookmark) * 2)
          branchClean[pruneSize * 2 + i + th_id] = branch[pruneSize * 2 + i + th_id] = sorted[wg_index + bookmark * 2 + i + th_id];        

      if(th_id == 0) {
        prune(branch, base, N, baseSize, pruneSize + branchSize - bookmark);
        
        j = 0;
        for (i = 0; i < pruneSize + branchSize - bookmark; i++)
            if (branch[i * 2] > -1) {
              result[result_index + (baseSize - baseStartSize) * 2 + j * 2] = branchClean[i * 2];
              result[result_index + (baseSize - baseStartSize) * 2 + j * 2 + 1] = branchClean[i * 2 + 1];
              j++;
            }
      }
    }
}
